/**
 * Copyright 2012 ibiblio
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */


package org.ibiblio.terasaur.bencode;
import java.io.InputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.List;
import java.util.Map;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.TreeSet;
import java.util.Iterator;
import java.io.UnsupportedEncodingException;

/**
 * Bencode encoder/decoder
 * See http://en.wikipedia.org/wiki/Bencode
 *
 * Bencode supports four data types, which are mapped to Java as:
 *  * integer     -> Long
 *  * byte string -> ByteArray
 *  * list        -> ArrayList
 *  * dictionary  -> HashMap
 *
 * Decoding returns a generic Object whose type depends on the decoded
 * data.  For complex encoded data, the decoded data will be a deep nested
 * structure.  Simple data is uncommon in real world usage.
 *
 * Encoding returns a ByteArray, which supports both string and binary data.
 * For convenience, integer encoding supports the Integer type in addition
 * to Long.
 *
 * @author john
 */
public class Bencode
{
    private static final int INTEGER_MARKER = (int)'i';
    private static final int LIST_MARKER = (int)'l';
    private static final int DICTIONARY_MARKER = (int)'d';
    private static final int BYTESTRING_LENGTH_SEPARATOR = (int)':';
    private static final int TERMINATOR = (int)'e';

    private static final int INTEGER_TYPE = 1;
    private static final int BYTESTRING_TYPE = 2;
    private static final int LIST_TYPE = 3;
    private static final int DICTIONARY_TYPE = 4;

    private boolean debug = false;

    public Bencode()
    {
    }

    public void setDebug(boolean debug)
    {
        this.debug = debug;
    }

    public Object decode(InputStream input) throws BencodeException, IOException, NumberFormatException
    {
        Object output = this._decode(input);
        // Extraneous data is not permitted
        int i = input.read();
        if (i != -1) {
            throw new BencodeException("Found extraneous encoded data");
        }
        return output;
    }

    public Object _decode(InputStream input) throws BencodeException, IOException, NumberFormatException
    {
        int type = this._getTypeForEncodedStream(input);
        Object output;

        switch (type) {
            case INTEGER_TYPE:
                output = this.decodeInteger(input);
                break;
            case BYTESTRING_TYPE:
                output = this.decodeString(input);
                break;
            case LIST_TYPE:
                output = this.decodeList(input);
                break;
            case DICTIONARY_TYPE:
                output = this.decodeDictionary(input);
                break;
            default:
                throw new BencodeException("Invalid encoded data type");
        }

        return output;
    }

    private int _getTypeForEncodedStream(InputStream input) throws BencodeException, IOException
    {
        input.mark(1);
        int i = input.read();
        if (i == -1) {
            throw new BencodeException("InputStream read error");
        }
        input.reset();
        return this._getTypeForEncodedMarker(i);
    }

    private int _getTypeForEncodedMarker(int i)
    {
        int type;
        switch (i) {
            case INTEGER_MARKER:
                type = INTEGER_TYPE;
                break;
            case LIST_MARKER:
                type = LIST_TYPE;
                break;
            case DICTIONARY_MARKER:
                type = DICTIONARY_TYPE;
                break;
            default:
                type = BYTESTRING_TYPE;
                break;
        }
        return type;
    }

    /*
    private Object _getObjectForEncodingType(int type) throws BencodeException
    {
        Object o = null;
        if (type == LIST_TYPE) {
            o = new ArrayList<Object>();
        } else if (type == DICTIONARY_TYPE) {
            o = new HashMap<String, Object>();
        } else if (type == BYTESTRING_TYPE) {
            o = new ByteArray();
        } else if (type == INTEGER_TYPE) {
            throw new BencodeException("Cannot decode an integer directly");
        }
        return o;
    }
    */

    /**
     * Decode an encoded integer value
     *
     * @param InputStream input
     * @return Long
     * @throws IOException
     * @throws NumberFormatException
     */
    private Long decodeInteger(InputStream input) throws IOException, NumberFormatException, BencodeException {
        StringBuffer buffer = new StringBuffer();
        boolean terminated = false;
        int i;
        char c;

        while (!terminated) {
            i = input.read();
            if (i == -1) {
                break; // no data from read
            }

            c = (char)i;

            if (this.debug) {
                System.out.println("char in decodeInteger: ]" + c + "[");
            }

            if (i == TERMINATOR) {
                terminated = true;
            // TODO: this should probably move outside the loop, check for 'i' as first element
            } else if (c != 'i') {
                buffer.append(c);
            }
        }

        if (!terminated) {
            throw new BencodeException("Encoded integer not terminated");
        }

        // Special case checks
        if (buffer.length() > 1) {
            char firstChar = buffer.charAt(0);
            char secondChar = buffer.charAt(1);

            if (firstChar == '-' && secondChar == '0') {
                throw new BencodeException("Invalid integer encoding: negative zero");
            }
            if (firstChar == '0') {
                throw new BencodeException("Invalid integer encoding: leading zero");
            }
        }

        String s = buffer.toString();
        Long l = Long.valueOf(s);
        return l;
    }

    private ByteArray decodeString(InputStream input) throws IOException, BencodeException {
        StringBuffer lengthString = new StringBuffer();
        int lengthValue = -1;
        ByteBuffer buffer = null;
        boolean foundSeparator = false;
        boolean stop = false;
        int byteCount = 0;
        int i;

        while (!stop) {
            i = input.read();
            if (i == -1) {
                break; // no data from read
            }

            // Allocate byte buffer as soon as we get a string length value
            if (buffer == null && lengthValue > 0) {
                buffer = ByteBuffer.allocate(lengthValue);
            }

            if (this.debug) {
                System.out.println("char in decodeString: ]" + (char)i + "[");
            }

            // Parse string length
            if (!foundSeparator) {
                foundSeparator = this._handleLengthCharacter(i, lengthString);
                if (foundSeparator) {
                    // Leading zero is not valid
                    if ((lengthString.length() > 1) && (lengthString.charAt(0) == '0')) {
                        throw new BencodeException("Invalid string length encoding: leading zero");
                    }
                    lengthValue = Integer.valueOf(lengthString.toString());
                }
            } else {
                buffer.put((byte)i);
                ++byteCount;
            }

            if (lengthValue > -1 && lengthValue <= byteCount) {
                stop = true;
            }
        }

        if (lengthValue == -1) {
            throw new BencodeException("Missing length value in string");
        }

        if (byteCount != lengthValue) {
            throw new BencodeException("String length mismatch");
        }

        ByteArray output = new ByteArray();
        if (buffer == null) {
            output.setBytes(new byte[0]);
        } else {
            output.setBytes(buffer.array());
        }
        return output;
    }

    /**
     * Process a single character in the input stream, looking for a string length
     * separator.  Append the character to the lengthString and return true if
     * we find the separator.
     *
     * @param i
     * @param lengthString
     * @return
     * @throws BencodeException
     */
    private boolean _handleLengthCharacter(int i, StringBuffer lengthString) throws BencodeException
    {
        char c = (char)i;
        boolean found = false;

        if (c == BYTESTRING_LENGTH_SEPARATOR) {
            found = true;
        } else if (i < 48 || i > 57) {
            throw new BencodeException("Invalid character found for string length");
        } else {
            lengthString.append(c);
        }

        return found;
    }

    private Object decodeList(InputStream input) throws IOException, BencodeException {
        if (!this._checkListMarker(input)) {
            return null;
        }

        ArrayList<Object> output = new ArrayList<Object>();
        int i;
        boolean terminated = false;

        while (!terminated) {
            input.mark(1);
            i = input.read();
            if (i == -1) {
                break; // no data from read
            }

            if (this.debug) {
                System.out.println("char in decodeList: ]" + (char)i + "[");
            }

            if (i == TERMINATOR) {
                terminated = true;
            } else {
                input.reset();
                this._addToList(input, i, output);
            }
        }

        if (!terminated) {
            throw new BencodeException("Encoded list not terminated");
        }

        return output;
    }

    private boolean _checkListMarker(InputStream input) throws IOException
    {
        // Check for valid list prefix
        int i = input.read();

        if (this.debug) {
            System.out.println("char in _checkListMarker: ]" + (char)i + "[");
        }

        if (i == -1 || (char)i != LIST_MARKER) {
            return false;
        } else {
            return true;
        }
    }

    private void _addToList(InputStream input, int encodingMarker, List<Object> outputList) throws BencodeException, IOException
    {
        int type = this._getTypeForEncodedMarker(encodingMarker);
        if (type == INTEGER_TYPE) {
            Long l = this.decodeInteger(input);
            outputList.add(l);
        } else {
            Object o = this._decode(input);
            outputList.add(o);
        }
    }

    private Object decodeDictionary(InputStream input) throws IOException, BencodeException
    {
        if (!this._checkDictionaryMarker(input)) {
            return null;
        }

        int i;
        boolean terminated = false;
        ByteArray previousKey = null;
        HashMap<String, Object> output = new HashMap<String, Object>();

        while (!terminated) {
            input.mark(1);
            i = input.read();
            if (i == -1) {
                break; // no data from read
            }

            if (this.debug) {
                System.out.println("char in decodeDictionary: ]" + (char)i + "[");
            }

            if (i == TERMINATOR) {
                terminated = true;
            } else {
                input.reset();
                ByteArray key = (ByteArray)this.decodeString(input);

                // Keys must be in lexicographical order
                if (previousKey != null) {
                    if (key.compareTo(previousKey) <= 0) {
                        throw new BencodeException("Dictionary keys not in order");
                    }
                }
                previousKey = key;
                Object value = this._decode(input);
                output.put(key.toString(), value);
            }
        }

        if (!terminated) {
            throw new BencodeException("Encoded list not terminated");
        }

        return (Map<String, Object>)output;
    }

    private boolean _checkDictionaryMarker(InputStream input) throws IOException
    {
        // Check for valid dictionary prefix
        int i = input.read();

        if (this.debug) {
            System.out.println("char in _checkDictionaryMarker: ]" + (char)i + "[");
        }

        if (i == -1 || (char)i != DICTIONARY_MARKER) {
            return false;
        } else {
            return true;
        }
    }

    public ByteArrayOutputStream encode(Object input) throws BencodeException, IOException
    {
        ByteArrayOutputStream output = new ByteArrayOutputStream();
        this._encode(input, output);
        return output;
    }

    private void _encode(Object input, ByteArrayOutputStream output) throws BencodeException, IOException
    {
        if (input instanceof Integer || input instanceof Long) {
            this._encodeInteger(input, output);
        } else if (input instanceof ByteArray) {
            this._encodeString(input, output);
        } else if (input instanceof ArrayList) {
            this._encodeList(input, output);
        } else if (input instanceof HashMap) {
            this._encodeDictionary(input, output);
        } else {
            throw new BencodeException("Cannot encode unsupported data type");
        }
    }

    private void _encodeInteger(Object input, ByteArrayOutputStream output) throws BencodeException
    {
        Long inputLong = (Long)input;
        output.write(INTEGER_MARKER);
        this._addLongToOutputStream(inputLong, output);
        output.write(TERMINATOR);
    }

    private void _addLongToOutputStream(Long inputLong, ByteArrayOutputStream output) throws BencodeException
    {
        byte[] bytes;
        try {
            bytes = inputLong.toString().getBytes("ISO-8859-1");
        } catch (UnsupportedEncodingException e) {
            throw new BencodeException(e.toString());
        }

        for (int i = 0; i < bytes.length; ++i) {
            output.write(bytes[i]);
        }
    }

    private void _encodeString(Object input, ByteArrayOutputStream output) throws BencodeException, IOException
    {
        ByteArray inputBytes = (ByteArray)input;
        Long length = new Long(inputBytes.getLength());
        this._addLongToOutputStream(length, output);
        output.write(BYTESTRING_LENGTH_SEPARATOR);
        output.write(inputBytes.getArray());
    }

    @SuppressWarnings("unchecked")
    private void _encodeList(Object input, ByteArrayOutputStream output) throws BencodeException, IOException
    {
        ArrayList<Object> inputList = (ArrayList<Object>)input;
        output.write(LIST_MARKER);
        for (Object item : inputList) {
            this._encode(item, output);
        }
        output.write(TERMINATOR);
    }

    @SuppressWarnings("unchecked")
    private void _encodeDictionary(Object input, ByteArrayOutputStream output) throws BencodeException, IOException
    {
        HashMap<String, Object> inputMap = (HashMap<String, Object>)input;
        TreeSet<String> keys = new TreeSet<String>(inputMap.keySet());
        Iterator<String> iter = keys.iterator();
        Object key;
        String keyString;

        output.write(DICTIONARY_MARKER);
        while (iter.hasNext()) {
            key = iter.next();
            if (!(key instanceof String)) {
                throw new BencodeException("Invalid dictionary key data type");
            }
            keyString = (String)key;
            this._encodeString(new ByteArray(keyString.getBytes()), output);
            Object item = inputMap.get(key);
            this._encode(item, output);
        }
        output.write(TERMINATOR);
    }
}
